import torch
from torch import nn
import time

from .build import ARCH_REGISTRY
from ..backbone.efficientdet_se import EfficientDetBackbone
from ..loss import FocalLoss
from utils import box_transform

@ARCH_REGISTRY.register()
class EfficientDetSE(nn.Module):
    def __init__(self, cfg):
        super(EfficientDetSE, self).__init__()
        self.criterion = FocalLoss(
            classification_loss_weight=cfg.MODEL.CLASSIFICATION_LOSS_WEIGHT,
            regression_loss_weight=cfg.MODEL.REGRESSION_LOSS_WEIGHT,
        )
        self.model = EfficientDetBackbone.from_pretrained(
            num_classes=cfg.MODEL.NUM_CLASSES,
            compound_coef=cfg.MODEL.COMPOUND_COEF,
            load_weights=cfg.MODEL.LOAD_PRETRAINED_WEIGHTS,
            in_channels=cfg.MODEL.INPUT.IN_CHANNELS,
            ratios=cfg.MODEL.RATIOS,
            scales=cfg.MODEL.SCALES,
            regressor_seblock_setting=cfg.MODEL.REGRESSOR_SEBLOCK,
            classifier_seblock_setting=cfg.MODEL.CLASSIFIER_SEBLOCK,
            efficientnet_seblock_setting=cfg.MODEL.EFFICIENTNET_SEBLOCK,
        )

    def forward(self, data):
        images = data['image']
        annotations = data.get('annotations', None)
        
        regression, classification, anchors = self.model(images)

        bboxes = box_transform(anchors, regression)
        loss_dict = {}
        if annotations is not None:
            cls_loss, reg_loss = self.criterion(classification, regression, anchors, annotations)
            loss_dict['cls_loss'] = cls_loss
            loss_dict['reg_loss'] = reg_loss
        return loss_dict, {"bboxes": bboxes, "classification": classification}
